# 解决图片扫描的pdf的文字提取问题
# 先将pdf转化为图片，从图片中提取文字

import os,shutil
import fitz
from aip import AipOcr

APP_ID = '25245042'
API_KEY = 'y1GoMvXLTGNKwrG2Aip5x6pE'
SECRET_KEY = 'vkdKjGNAo9sFKsT2WEqmklGp3yEZGHUT'

client = AipOcr(APP_ID, API_KEY, SECRET_KEY)


src_file = './PDF/normal'  # 源文件，在normal中有但是txt中没有的文件
tar_file = './PDF_PIC'  # 将normal中有但是txt没有的文件提取出来，放在tar_file下

imagePath = './temp/' #当前的pdf文件转成的img的位置


# # 找出没有被提取出来的pdf
# def find():
#     list = os.listdir(src_file)
#     for i in list:
#         name = os.path.splitext(i)[0]
#         if (os.path.exists("./PDF/txt/" + name + '.txt')):
#             continue
#         else:
#             shutil.move(src_file + '/' + i, tar_file)
#             print(i)


# 将该类pdf转换成图片
def toIMG(pdf_file):
    try:
        # 每次处理一个新的文件先清空临时存放图片的文件夹
        shutil.rmtree(imagePath)

        pdfDoc = fitz.open(pdf_file)
        for pg in range(pdfDoc.pageCount):
            page = pdfDoc[pg]
            rotate = int(0)
            # 每个尺寸的缩放系数为1.3，这将为我们生成分辨率提高2.6的图像。
            # 此处若是不做设置，默认图片大小为：792X612, dpi=96
            zoom_x = 1.33333333  # (1.33333333-->1056x816)   (2-->1584x1224)
            zoom_y = 1.33333333
            mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
            pix = page.get_pixmap(matrix=mat, alpha=False)

            if not os.path.exists(imagePath):  # 判断存放图片的文件夹是否存在
                os.makedirs(imagePath)  # 若图片文件夹不存在就创建
            pix.save(imagePath + '/' + 'images_%s.png' % pg)  # 将图片写入指定的文件夹内
    except BaseException as e:
        with open('./yan/log.txt', "a") as file:  # "a"代表追加内容
            file.write(pdf + " " + repr(e) + "\n")

# 从图片中提取文字
def fun(pdf):
    try:
        list = os.listdir(imagePath)

        for img in list:
            image = open(imagePath + img, 'rb').read()
            options = {}
            options["detect_direction"] = "true"

            result = client.basicGeneral(image, options)
            words = result['words_result']

            txt = './yan/' + os.path.splitext(pdf)[0] + '.txt'
            with open(txt, "a") as file:  # 只需要将之前的”w"改为“a"即可，代表追加内容
                for line in words:
                    file.write(line['words'] + "\n")
    except BaseException as e:
        with open('./yan/log.txt', "a") as file:  # "a"代表追加内容
            file.write(pdf + " " + repr(e) + "\n")


if __name__ == '__main__':
    list = os.listdir(tar_file)
    for pdf in list:
        print(pdf)
        pdf_file = './PDF_PIC/' + pdf
        toIMG(pdf_file)
        fun(pdf)


